import sympy as sym
import numpy as np


def has(expr, subexpr):
    """
    has(expr,subexpr) returns logical 1 (true) if expr contains subexpr. 
    Otherwise, it returns logical 0 (false).
    If expr is an array, has(expr,subexpr) returns an array of the same size as expr. 
    The returned array contains logical 1s (true) where the elements of expr contain subexpr,
    and logical 0s (false) where they do not.
    If subexpr is an array, has(expr,subexpr) checks if expr contains any element of subexpr.

    Args:
        expr : symbolic expression | symbolic expression array 
        subexpr : symbolic expression | symbolic expression array 

    Returns:
        the result of a and b(a & b).
    """
    # Case1: expr and subexpr are simple expression
    if isinstance(expr, sym.Expr) and isinstance(subexpr, sym.Expr):
        # print("case1")
        return expr.has(subexpr)
    elif isinstance(expr, (list, np.ndarray)) and isinstance(subexpr, sym.Expr):
        # print("case2")
        narr = np.array(expr)
        result = np.full(narr.shape, False)
        if len(narr.shape) == 1:
            for idx in range(narr.size):
                result[idx] = has(narr[idx], subexpr)
        elif len(narr.shape) == 2:
            row = narr.shape[0]
            col = narr.shape[1]
            for r in range(row):
                for c in range(col):
                    result[r][c] = has(narr[r][c], subexpr)
        else:
            raise ValueError("Array has too many dimensions")
        return result
    elif isinstance(expr, sym.Expr) and isinstance(subexpr, (list, np.ndarray)):
        narr = np.array(subexpr)
        result = False
        if len(narr.shape) == 1:
            for idx in range(narr.size):
                result |= has(expr, narr[idx])
        elif len(narr.shape) == 2:
            row = narr.shape[0]
            col = narr.shape[1]
            for r in range(row):
                for c in range(col):
                    result |= has(expr, narr[r][c])
        else:
            raise ValueError("Array has too many dimensions")
        return result
    elif isinstance(expr, (list, np.ndarray)) and isinstance(subexpr, (list, np.ndarray)):
        # print("case4")
        narr = np.array(expr)
        result = np.full(narr.shape, False)
        if len(narr.shape) == 1:
            for idx in range(narr.size):
                result[idx] = has(narr[idx], subexpr)
        elif len(narr.shape) == 2:
            row = narr.shape[0]
            col = narr.shape[1]
            for r in range(row):
                for c in range(col):
                    result[r][c] = has(narr[r][c], subexpr)
        else:
            raise ValueError("Array has too many dimensions")
        return result
    else:
        raise TypeError("Input type not supported.")


x, y, z = sym.symbols('x y z')
# Test Case 1: expr and subexpr are all simple expression
expr1 = x + y + z + 3
expr2 = y
assert has(expr1, expr2) == True

# Test Case 2: expr is a expression array, subexpr is simple expression
expr1 = [
    [x + y + z],
    [3*x + y]
]
expr2 = y + z
result = np.array([
    [True],
    [False]
])
assert np.array_equal(has(expr1, expr2),  result)

# Test Case 3: expr is simple expression, subexpr is a expression array
expr1 = x + y + sym.cos(z)
expr2 = [
    [sym.sin(x), sym.cos(y)],
    [y, z]
]
assert has(expr1, expr2) == True

# Test Case 4: expr is simple expression, subexpr is a expression array
expr1 = [
    [x + y + z, sym.log(x)+sym.log(y)],
    [x**2+y**3, sym.exp(z) - sym.log(x)]
]
expr2 = sym.log(x)
result = np.array([
    [False, True],
    [False, True]
])
assert np.array_equal(has(expr1, expr2),  result)

# Test Case 5: expr and subexpr are all expression array
expr1 = [
    [x + y + z, sym.log(x)+sym.log(y)],
    [x**2+y**3, sym.exp(z) - sym.log(x)]
]
expr2 = [sym.exp(z), y**3]
result = np.array([
    [False, False],
    [True, True]
])
assert np.array_equal(has(expr1, expr2),  result)
